import binascii
import rsa
from Crypto.Cipher import AES
from binascii import b2a_hex, a2b_hex


# import base64
# from Crypto import Random


# 定义秘钥长度
KEY_LENGTH = 512


# 获取 RSA 秘钥
# return <rsa key object>
def get_key():
    # 生成公私钥
    _public_key, _private_key = rsa.newkeys(KEY_LENGTH)

    return _public_key, _private_key


# 获取给定 RSA 公私钥的字符串形式
# <rsa key object> -> <str>
# return <str>
def get_key_str_from_obj(_public_key, _private_key):

    _pub_key = bytes.decode(_public_key.save_pkcs1())
    _pri_key = bytes.decode(_private_key.save_pkcs1())

    return _pub_key, _pri_key


# 获取给定字符串形式 RSA 公私钥的 object 形式
# <str> -> <rsa key object>
# return <rsa key object>
def get_key_obj_from_str(_pub_key, _pri_key):
    _public_key = rsa.PublicKey.load_pkcs1(str.encode(_pub_key))
    _private_key = rsa.PrivateKey.load_pkcs1(str.encode(_pri_key))
    return _public_key, _private_key


def get_pub_key_obj_from_str(_pub_key):
    return rsa.PublicKey.load_pkcs1(str.encode(_pub_key))


# RSA 加密
# @attr: _original_text <str>: 代价密的明文
#        _public_key <rsa public key object>: RSA 公钥
# return: <str> 加密后的密文
def encrypt_rsa(_original_text, _public_key):
    # 原文重新编码
    _plaintext = _original_text.encode("utf-8")
    # 密文
    _ciphertext = rsa.encrypt(_plaintext, _public_key)
    # 返回tuple
    return bytes.decode(binascii.hexlify(_ciphertext))


# RSA 解密
# @attr: _ciphertext <str>: 加密后的密文
#        _private_key <rsa private key object>: RSA 私钥
# return: <str> 解密后的明文
def decrypt_rsa(_ciphertext, _private_key):
    _ciphertext = binascii.unhexlify(str.encode(_ciphertext))
    # 解密
    _plaintext = rsa.decrypt(_ciphertext, _private_key)
    # 解码
    _original_text = _plaintext.decode('utf-8')
    return _original_text


# AES-CBC 加密
# 如果text不足16位的倍数就用空格补足为16位
def add_to_16(text):
    if len(text.encode('utf-8')) % 16:
        add = 16 - (len(text.encode('utf-8')) % 16)
    else:
        add = 0
    text = text + ('\0' * add)
    return text.encode('utf-8')


# AES 加密
# @attr: _text <str>: 待价密的明文
#        _key <str>: AES （对称）秘钥 16位
# return: <str> 加密后的密文
def encrypt_aes(text, key):
    key = key.encode('utf-8')
    mode = AES.MODE_CBC
    iv = b'qqqqqqqqqqqqqqqq'
    text = add_to_16(text)
    cryptos = AES.new(key, mode, iv)
    cipher_text = cryptos.encrypt(text)
    # AES加密后的字符串不一定是ascii字符集 转为16进制字符串
    return bytes.decode(b2a_hex(cipher_text))


# AES 解密
# @attr: _text <str>: 待解密的密文
#        _key <str>: AES （对称）秘钥 16位
# return: <str> 解密后的明文
def decrypt_aes(text, key):
    text = str.encode(text)
    key = key.encode('utf-8')
    iv = b'qqqqqqqqqqqqqqqq'
    mode = AES.MODE_CBC
    cryptos = AES.new(key, mode, iv)
    plain_text = cryptos.decrypt(a2b_hex(text))
    return bytes.decode(plain_text).rstrip('\0')

